import os
import importlib, sys
import requests
import datetime
from Crypto.Cipher import AES
from binascii import b2a_hex, a2b_hex

from demo import config

importlib.reload(sys)
timeout=5


# sys.setdefaultencoding('utf-8')


def download(url, download_path, file_name):
    all_content = requests.get(url, timeout=timeout).text  # 获取第一层M3U8文件内容
    if "#EXTM3U" not in all_content:
        raise BaseException("非M3U8的链接")

    if "EXT-X-STREAM-INF" in all_content:  # 第一层
        file_line = all_content.split("\n")
        for line in file_line:
            if '.m3u8' in line:
                url = url.rsplit("/", 1)[0] + "/" + line  # 拼出第二层m3u8的URL
                all_content = requests.get(url, timeout=timeout).text

    file_line = all_content.split("\n")

    unknow = True
    key = ""
    for index, line in enumerate(file_line):  # 第二层
        if "#EXT-X-KEY" in line:  # 找解密Key
            method_pos = line.find("METHOD")
            comma_pos = line.find(",")
            method = line[method_pos:comma_pos].split('=')[1]
            print
            "Decode Method：", method

            uri_pos = line.find("URI")
            quotation_mark_pos = line.rfind('"')
            key_path = line[uri_pos:quotation_mark_pos].split('"')[1]

            key_url = url.rsplit("/", 1)[0] + "/" + key_path  # 拼出key解密密钥URL
            res = requests.get(key_url, timeout=timeout)
            key = res.content
            print
            "key：", key

        if "EXTINF" in line:  # 找ts地址并下载
            unknow = False
            pd_url = url.rsplit("/", 1)[0] + "/" + file_line[index + 1]  # 拼出ts片段的URL
            # print pd_url

            res = requests.get(pd_url, timeout=timeout)
            c_fule_name = file_line[index + 1].rsplit("/", 1)[-1]

            if len(key):  # AES 解密
                cryptor = AES.new(key, AES.MODE_CBC, key)
                with open(os.path.join(download_path, c_fule_name + ".mp4"), 'ab') as f:
                    f.write(cryptor.decrypt(res.content))
            else:
                with open(os.path.join(download_path, c_fule_name), 'ab') as f:
                    f.write(res.content)
                    f.flush()
    if unknow:
        raise BaseException("未找到对应的下载链接")
    else:
        print(f"下载完成，保存路径[{download_path}/{file_name}]")
    merge_file(download_path, file_name)


def merge_file(path, file_name):
    # ts文件绝对路径
    ts_path = path
    # 读取ts文件夹下所有的ts文件
    path_list = os.listdir(ts_path)
    # 对文件进行排序
    path_list.sort()
    # 将排序后的ts的绝对路径放入列表中
    li = [os.path.join(ts_path, filename) for filename in path_list]
    # 类似于[001.ts|00.2ts|003.ts]
    input_file = '|'.join(li)
    # 指定输出文件名称
    output_file = file_name
    # 使用ffmpeg将ts合并为mp4
    command = f'{config.ffmpeg_path} -i "concat:%s" -acodec copy -vcodec copy -absf aac_adtstoasc %s -loglevel quiet' % (
        input_file, output_file)
    # 指行命令
    os.chdir(path)
    os.system(command)
    os.system('rm -rf *.ts')
    # os.chdir(path)
    # cmd = "cat * > new.tmp"
    # os.system(cmd)
    # os.rename("new.tmp", "new.mp4")

    # 回到项目根目录
    os.chdir(os.path.abspath('../../'))


if __name__ == '__main__':
    url = "http://zblj-w.hljtv.com:8080/live/qiqihaerdongwuyoueryuandongbeihu.m3u8"
    # 获得download文件夹
    download_path = os.getcwd() + "/download"
    if not os.path.exists(download_path):
        os.mkdir(download_path)

    # 新建日期文件夹
    download_path = os.path.join(download_path, datetime.datetime.now().strftime('%Y%m%d_%H%M%S'))
    os.mkdir(download_path)

    file_name = 'new.mp4'

    download(url, download_path, file_name)
